import sys

sys.path += ['./']

from utils.io import read_pkl, write_pkl
import time
import faiss
import os
import numpy as np
import argparse


def faiss_retrieval(querys, dataset, write_path, k=1, shard=True):
    dim = 1024
    # cpu_index = faiss.index_factory(dim, 'OPQ32,IVF1024,PQ32')
    cpu_index = faiss.IndexFlatIP(dim)
    ngpus = faiss.get_num_gpus()
    print("number of GPUs:", ngpus)
    co = faiss.GpuMultipleClonerOptions()
    co.shard = shard
    gpu_index = faiss.index_cpu_to_all_gpus(cpu_index, co=co)

    print('now load query')
    xq = np.fromfile(querys, dtype=np.float32, count=-1)
    xq.resize(xq.shape[0] // dim, dim)
    print('now load dataset')
    xb = np.fromfile(dataset, dtype=np.float32, count=-1)
    xb.resize(xb.shape[0] // dim, dim)

    xb = np.array(xb).astype('float32')
    xq = np.array(xq).astype('float32')
    print('documents=', len(xb), ', queries=', len(xq))

    # gpu_index.train(xb)
    gpu_index.add(xb)
    print(gpu_index.ntotal)

    t = time.time()
    D, I = gpu_index.search(xq, k)
    print('time cost =', time.time() - t)

    write_pkl(D, write_path + '_D.pkl')
    write_pkl(I, write_path + '_I.pkl')


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-q', type=str)
    parser.add_argument('-d', type=str)
    parser.add_argument('-o', type=str)
    parser.add_argument('--rows', type=int, default=1)
    args = parser.parse_args()

    query_path = args.q
    document_path = args.d
    output_path = args.o
    rows = args.rows

    print('FAISS retrieval')

    faiss_retrieval(query_path, document_path, output_path, k=rows)


if __name__ == '__main__':
    main()
